# -*- coding: utf-8 -*-
# ----------------------------
# @Time    : 2022/6/4 2:35 PM
# @Author  : changqingai
# @FileName: tf_device.py
# ----------------------------

import tensorflow as tf

physical_devices = tf.config.list_physical_devices('CPU')
print("physical_devices: ", physical_devices)
assert len(physical_devices) == 1, "No CPUs found"
configs = tf.config.get_logical_device_configuration(
  physical_devices[0])
try:
  assert configs is None
  tf.config.set_logical_device_configuration(
    physical_devices[0],
    [tf.config.LogicalDeviceConfiguration(),
     tf.config.LogicalDeviceConfiguration()])
  configs = tf.config.get_logical_device_configuration(
    physical_devices[0])
  assert len(configs) == 2
except:
  # Cannot modify virtual devices once initialized.
  pass

logical_device = tf.config.get_logical_device_configuration(physical_devices[0])
print("logical_device:", logical_device)
